﻿using System;
using System.Collections;
using System.Collections.Generic;
using System.Linq;
using System.Runtime.InteropServices;
using Unity.Collections;
using UnityEngine;
using UnityEngine.Rendering;

namespace MarchingSquares
{
    public class MarchingGPU : Marching
    {
        // Compute shader
        private readonly ComputeShader _marchingShader;

        // Compute shaders variables
        private ComputeBuffer _voxelDataBuffer;
        private ComputeBuffer _trianglesBuffer;
        private ComputeBuffer _verticesBuffer;
        private ComputeBuffer _uvsBuffer;
        
        private int[] _trianglesArray;
        private Vector3[] _verticesArray;
        private Vector3[] _uvsArray;

        // Parameters and settings
        private int _resolution;

        // Thread params
        private int _threadGroups;
        private const int ThreadPerGroups = 128;

        // Chunk Material and Position
        private Material _material;
        private Vector3 _position;
        private static readonly int VoxelsBuffer = Shader.PropertyToID("voxels");
        private ComputeBuffer argsBuffer;
        private uint[] args = new uint[5] { 0, 0, 0, 0, 0 };

        public MarchingGPU(ComputeShader marching)
        {
            _marchingShader = marching;
        }

        public override void Initialize(Material voxelMaterial, Vector3 chunkPosition, int chunkResolution,
            VoxelChunk voxelChunk)
        {
            _voxelChunk = voxelChunk;
            // Create the base mesh
            mesh = new Mesh
            {
                indexFormat = IndexFormat.UInt32,
                name = "VoxelGrid Mesh"
            };

            // Assign the material and the position
            _material = voxelMaterial;
            _position = chunkPosition;

            // Set the chunk resolution
            _resolution = chunkResolution;

            // Setup the thread groups
            _threadGroups = Mathf.CeilToInt(_resolution * _resolution / (float)ThreadPerGroups);

            argsBuffer = new ComputeBuffer(1, args.Length * sizeof(uint), ComputeBufferType.IndirectArguments);
            Refresh();
        }

        public void Refresh()
        {
            _trianglesBuffer?.Release();
            _verticesBuffer?.Release();
            _voxelDataBuffer?.Release();
            _uvsBuffer?.Release();
            
            int voxelNumber = _resolution * _resolution;

            // Setup the output arrays
            _trianglesArray = new int[voxelNumber * 9];
            _verticesArray = new Vector3[voxelNumber * 6];
            _uvsArray = new Vector3[voxelNumber * 6];

            // Calculate the compute buffer again Marshal.SizeOf(typeof(Vector3)) + Marshal.SizeOf(typeof(float)));
            _voxelDataBuffer = new ComputeBuffer(voxelNumber,
                Marshal.SizeOf(typeof(VoxelData)));
            _trianglesBuffer = new ComputeBuffer(voxelNumber * 9, Marshal.SizeOf(typeof(int)));
            _verticesBuffer = new ComputeBuffer(voxelNumber * 6, Marshal.SizeOf(typeof(Vector3)));
            _uvsBuffer = new ComputeBuffer(voxelNumber * 6, Marshal.SizeOf(typeof(Vector3)));
            
        }

        public override void Triangulate(ref VoxelData[] voxelData, float isoLevel, bool useUVMapping = false,
            bool useInterpolation = true,int[] cellsType=null)
        {
            if (_waitDataQueue.Any())
            {
                Graphics.DrawMesh(mesh, _position, Quaternion.identity, _material, 0);
                return;
            }

            // Clear the mesh
            Clear();
            // Compress the mesh
            // Indirect args
            var subMeshIndex = 0;
            if (mesh != null)
            {
                // 0 == number of triangle indices, 1 == population, others are only relevant if drawing submeshes.
                args[0] = (uint)mesh.GetIndexCount(subMeshIndex);
                args[1] = (uint)(_resolution * _resolution);
                args[2] = (uint)mesh.GetIndexStart(subMeshIndex);
                args[3] = (uint)mesh.GetBaseVertex(subMeshIndex);
            }
            else
            {
                args[0] = args[1] = args[2] = args[3] = 0;
            }

            argsBuffer.SetData(args);
            // List<Vector3> verticesList = new List<Vector3>();
            // List<int> trianglesList = new List<int>();
            //
            // CompressMesh(ref _trianglesArray, ref _verticesArray, ref trianglesList, ref verticesList);

            // Set the mesh vertices and triangles
            mesh.SetVertices(_verticesArray);
            mesh.SetTriangles(_trianglesArray, 0);

            if (useUVMapping)
            {
                mesh.SetUVs(0, _uvsArray);
            }
            //
            // // Mark mesh as dynamic
            // mesh.MarkDynamic(); 


            // Draw the mesh GPU
            Graphics.DrawMesh(mesh, _position, Quaternion.identity, _material, 0);
            // var bound = new Bounds(Vector3.zero, new Vector3(100.0f, 100.0f, 100.0f));
            // Graphics.DrawMeshInstancedIndirect(mesh,subMeshIndex,_material,bound,argsBuffer);
            // Compute
            int marchingKernel = _marchingShader.FindKernel("March");
            int cells = _resolution - 1;

            Array.Clear(_trianglesArray, 0, _trianglesArray.Length - 1);
            Array.Clear(_verticesArray, 0, _verticesArray.Length - 1);
            Array.Clear(_uvsArray, 0, _uvsArray.Length - 1);
            // Push our prepared data into Buffer
            _voxelDataBuffer.SetData(voxelData);
            _trianglesBuffer.SetData(_trianglesArray);
            _verticesBuffer.SetData(_verticesArray);
            _uvsBuffer.SetData(_uvsArray);

            
            // Set the buffers
            _marchingShader.SetBuffer(marchingKernel, "voxels", _voxelDataBuffer);
            _marchingShader.SetBuffer(marchingKernel, "triangles", _trianglesBuffer);
            _marchingShader.SetBuffer(marchingKernel, "vertices", _verticesBuffer);
            _marchingShader.SetBuffer(marchingKernel, "uvs", _uvsBuffer);

            // Update the parameters data and dispatch it
            _marchingShader.SetInt("cells", cells);
            _marchingShader.SetFloat("isoLevel", isoLevel);
            _marchingShader.SetBool("useInterpolation", useInterpolation);
            _marchingShader.SetBool("useUVMapping", useUVMapping);

            _material.SetBuffer(VoxelsBuffer, _voxelDataBuffer);


            _marchingShader.Dispatch(marchingKernel, _threadGroups, 1, 1);

            // Get the data from the buffers
            // _trianglesBuffer.GetData(_trianglesArray);
            // _verticesBuffer.GetData(_verticesArray);
            // _uvsBuffer.GetData(_uvsArray);

            //Async GetData
            // _voxelChunk.StartCoroutine(GetDataAsync());
            GetDataAsync();
            // _voxelDataBuffer.GetData(voxelData);
        }

        public override void Clear() => mesh.Clear();

        public override void Destroy()
        {
            // Dispose the compute buffers
            _voxelDataBuffer?.Dispose();
            _trianglesBuffer?.Dispose();
            _verticesBuffer?.Dispose();
            _uvsBuffer?.Dispose();

            _voxelDataBuffer = null;
            _trianglesBuffer = null;
            _verticesBuffer = null;
            _uvsBuffer = null;

            argsBuffer?.Release();
            argsBuffer = null;
        }

        private Queue<int> _waitDataQueue = new Queue<int>();

        void GetDataAsync()
        {
            // IEnumerator AsyncGetData(ComputeBuffer buffer,Array data)
            // {
            //     var req = AsyncGPUReadback.Request(buffer);
            //     yield return new WaitUntil(() => req.done);
            //     var data = req.GetData<Color32>().ToArray();
            // }
            void OnCompleteReadback(AsyncGPUReadbackRequest request)
            {
                if (!request.done)
                {
                    Debug.Log("readback hasnt done yet");
                    return;
                }

                if (request.hasError)
                {
                    Debug.Log("readback error");
                }
                else
                {
                    // GetDataFromGPU(request);
                }
            }

            var req1 = AsyncGPUReadback.Request(_trianglesBuffer, req =>
            {
                // Array.Clear(_trianglesArray, 0, _trianglesArray.Length - 1);
                req.GetData<int>().CopyTo(_trianglesArray);
                _waitDataQueue.Dequeue();
            });
            var req2 = AsyncGPUReadback.Request(_verticesBuffer, req =>
            {
                // Array.Clear(_verticesArray, 0, _verticesArray.Length - 1);
                req.GetData<Vector3>().CopyTo(_verticesArray);
                _waitDataQueue.Dequeue();
            });
            var req3 = AsyncGPUReadback.Request(_uvsBuffer, req =>
            {
                // Array.Clear(_uvsArray, 0, _uvsArray.Length - 1);
                req.GetData<Vector3>().CopyTo(_uvsArray);
                _waitDataQueue.Dequeue();
            });
            _waitDataQueue.Enqueue(1);
            _waitDataQueue.Enqueue(1);
            _waitDataQueue.Enqueue(1);
            // yield return new WaitUntil(() => req1.done);
            // yield return new WaitUntil(() => req2.done);
            // yield return new WaitUntil(() => req3.done);

            // _trianglesArray = req1.GetData<int>().ToArray();
            // _verticesArray = req2.GetData<Vector3>().ToArray();
            // _uvsArray = req3.GetData<Vector3>().ToArray();
        }
        
    }
}